import random
from dataclasses import dataclass, field
from typing import List, Tuple, Optional
import matplotlib.pyplot as plt

# 定义颜色常量（使用RGB元组）
TEACHER_SEAT_COLOR = (255, 255, 255)  # 白色
OBSTACLES_COLOR = (126, 126, 126)    # 灰色
TURQUOISE = (64, 224, 208)           # 青绿色
BLACK = (0, 0, 0)
BLUE = (0, 0, 255)
GREEN = (0, 255, 0)
RED = (255, 0, 0)
BROWN = (165, 42, 42)
GRAY = (128, 128, 128)
YELLOW = (255, 255, 0)

@dataclass
class Patch:
    x: int
    y: int
    pcolor: Tuple[int, int, int]
    object_type: str
    energy: int = 0
    is_jumpable: bool = True
    outside: bool = False

    def neighbors(self, patches: List['Patch'], width: int, height: int) -> List['Patch']:
        neighbors = []
        directions = [(-1,0),(1,0),(0,-1),(0,1),(-1,-1),(-1,1),(1,-1),(1,1)]
        for dx, dy in directions:
            nx, ny = self.x + dx, self.y + dy
            if 0 <= nx < width and 0 <= ny < height:
                neighbors.append(patches[nx + ny * width])
        return neighbors

@dataclass
class Turtle:
    x: int
    y: int
    category: str
    speed: float
    current_speed: float = 0.0
    travelled_distance: float = 0.0
    can_jump_over_obstacles: bool = False
    will_die: bool = False
    is_panicked: bool = False
    my_teammate_id: Optional[int] = None
    inspire_students: bool = False

    def move(self, patches: List[Patch], width: int, height: int):
        # 获取当前位置的patch
        current_patch = patches[self.x + self.y * width]
        # 检查是否到达出口
        if current_patch.object_type == "door":
            evacuation_distances.append(self.travelled_distance)
            evacuation_durations.append(ticks)
            average_evacuation_distance = sum(evacuation_distances) / len(evacuation_distances)
            total_evacuation_distance = sum(evacuation_distances)
            average_evacuation_duration = sum(evacuation_durations) / len(evacuation_durations)
            # 人员离开模拟
            return False
        else:
            # 获取可移动的邻居patch
            possible_patches = [
                p for p in current_patch.neighbors(patches, width, height)
                if p.object_type in ["floor", "door", "intermediate-door"] and not p.outside
            ]
            if not possible_patches:
                return True  # 无法移动，保持当前位置
            # 选择能量最低的patch作为下一个移动目标
            next_patch = min(possible_patches, key=lambda p: p.energy)
            # 移动到下一个patch
            self.x, self.y = next_patch.x, next_patch.y
            self.travelled_distance += self.current_speed
            # 如果移动到障碍物，处理跳跃逻辑
            if next_patch.object_type == "obstacle":
                if self.can_jump_over_obstacles:
                    self.travelled_distance += 0.1 + random.uniform(0, self.current_speed)
                else:
                    self.travelled_distance += self.current_speed / 2
            # 更新恐慌状态
            if self.travelled_distance > 150:
                self.is_panicked = True
            # 如果有队友（残疾人），帮助其移动
            if self.my_teammate_id is not None:
                teammate = turtles[self.my_teammate_id]
                if teammate.will_die:
                    self.speed = teammate.current_speed
                    self.current_speed = self.speed
                self.x, self.y = teammate.x, teammate.y
            # 如果附近有老师，增加速度
            if any(p.object_type == "teacher" for p in current_patch.neighbors(patches, width, height)):
                self.current_speed += 0.1
            else:
                self.current_speed = self.speed
            return True

def setup_patches(width: int, height: int) -> List[Patch]:
    patches = []
    for y in range(height):
        for x in range(width):
            p = Patch(x, y, pcolor=BLACK, object_type="wall")
            patches.append(p)
    return patches

def setup_turtles(patches: List[Patch], width: int, height: int, number_of_teachers: int, number_of_children: int, percentage_children_with_disabilities: float) -> List[Turtle]:
    turtles = []
    # 放置教师
    for _ in range(number_of_teachers):
        while True:
            x, y = random.randint(0, width-1), random.randint(0, height-1)
            p = patches[x + y * width]
            if p.object_type == "teacher_seat" and not any(t.x == x and t.y == y for t in turtles):
                break
        t = Turtle(x, y, category="teacher", speed=1.2, can_jump_over_obstacles=False)
        turtles.append(t)
    # 放置学生
    for _ in range(number_of_children):
        while True:
            x, y = random.randint(0, width-1), random.randint(0, height-1)
            p = patches[x + y * width]
            if p.object_type == "seat" and not any(t.x == x and t.y == y for t in turtles):
                break
        category = "teenager"
        color = YELLOW
        speed = 1.02
        if random.uniform(0, 100) < percentage_children_with_disabilities:
            category = "pers_with_disabilities"
            color = RED
            speed = 0.42
        t = Turtle(x, y, category=category, speed=speed, can_jump_over_obstacles=False)
        turtles.append(t)
    return turtles

def setup_obstacles(patches: List[Patch], width: int, height: int, percentage_obstacles: float):
    for p in patches:
        if p.object_type == "floor" and random.uniform(0, 100) < percentage_obstacles:
            p.object_type = "obstacle"
            p.pcolor = OBSTACLES_COLOR
            p.is_jumpable = random.choice([True, False])

def setup_exits(patches: List[Patch], width: int, height: int):
    # 示例：设置一个出口在底部中央
    for x in range(width):
        p = patches[x + (height-1) * width]
        if x == width//2 or x == width//2 -1:
            p.object_type = "door"
            p.pcolor = TURQUOISE
        else:
            p.object_type = "wall"
            p.pcolor = BLACK

def simulate():
    width, height = 100, 50  # 环境尺寸
    patches = setup_patches(width, height)
    setup_exits(patches, width, height)
    number_of_teachers = 10
    number_of_children = 200
    percentage_children_with_disabilities = 5.0  # 5%
    turtles = setup_turtles(patches, width, height, number_of_teachers, number_of_children, percentage_children_with_disabilities)
    setup_obstacles(patches, width, height, percentage_obstacles=10.0)  # 10%障碍物

    global ticks, evacuation_distances, evacuation_durations, average_evacuation_distance, total_evacuation_distance, average_evacuation_duration
    ticks = 0
    evacuation_distances = []
    evacuation_durations = []
    average_evacuation_distance = 0.0
    total_evacuation_distance = 0.0
    average_evacuation_duration = 0.0

    while True:
        if not any(t.will_die for t in turtles):
            ticks += 1
            for t in turtles:
                if not t.move(patches, width, height):
                    t.will_die = True
            if ticks > 1000:
                break
        else:
            break

    print(f"疏散完成，总时间步: {ticks}")
    print(f"平均疏散距离: {average_evacuation_distance}")
    print(f"总疏散距离: {total_evacuation_distance}")
    print(f"平均疏散时间: {average_evacuation_duration}")

if __name__ == "__main__":
    simulate()